#Ensure you are using a GPU (otherwise change execution mode to GPU)
#You can also verify you have a powerful GPU (P4 < K80 < T4 < P100)
!nvidia-smi -L
import os
from os.path import *
basedir = "/content"
# Import library
repodir = join(basedir,"open-image-restoration")
if not exists(repodir):
os.chdir(basedir)
!git clone https://github.com/titsitits/open-image-restoration {repodir}
os.chdir(repodir)
#Todo: create a pip module
#Specific dependencies needed for colorization
!pip install --quiet -r requirements.txt
import ImagePipeline_utils as IP
import ImageRestorer
restorer = ImageRestorer.ImageRestorer()
#Import some files
inputdir = join(basedir,"input")
IP.initdir(inputdir)
outputdir = join(basedir,"output")
IP.initdir(outputdir)
import_local_files = False
from google.colab import files
os.chdir(inputdir)
if import_local_files:
files.upload()
else:
!wget -q -O "1.jpg" "https://upload.wikimedia.org/wikipedia/commons/b/bf/Christmas_Truce_by_Frederic_Villiers.jpg"
!wget -q -O "2.jpg" "https://upload.wikimedia.org/wikipedia/commons/6/60/Christmas_Truce_3.jpg"
!wget -q -O "3.jpg" "https://upload.wikimedia.org/wikipedia/commons/2/20/Weihachtsfriede.jpg"
!wget -q -O "4.jpg" "https://upload.wikimedia.org/wikipedia/commons/c/ce/The_Christmas_Truce_on_the_Western_Front%2C_1914_Q50721.jpg"
os.chdir(basedir)
#Todo: make restorer available from outside repo directory
os.chdir(repodir)
#Todo: improve log handling
raising = False
# Restore images
restorer.preprocess(inputdir, outputdir, gray=True)
restorer.remove_stripes(raising=raising, process_args="-n 2") #reduce image moire (n iterations)
restorer.denoise(raising=raising) #remove image grain
restorer.remove_stripes(raising=raising, process_args="-n 2") #reduce remaining image moire (n iterations)
restorer.colorize(raising=raising) #colorize image (first time takes a long time as large models must be downloaded)
restorer.super_resolution(raising=raising) #upsample image
IP.compare_folders([inputdir, outputdir], imagewidth=300, sameheight = False)
IP.compare_folders([inputdir, outputdir], imagewidth=1000, sameheight = False)